# -*- coding: utf-8 -*-
"""
Created on Mon Dec 25 13:56:27 2023

@author: HD
"""

import numpy as np
from statsmodels.tsa.arima.model import ARIMA


def arima_forecast(data, plan, f_plan, use_plan, order=(1, 1, 1), forecast_steps=3):
    """
    使用ARIMA模型对多个时间序列进行预测。

    参数：
    - data: 输入的时间序列数据，维度为（n_timesteps, n_samples）。
    - plan: 输入的发货计划序列数据，维度为（n_timesteps, n_samples），要求与data维度一样。
    - use_plan: 是否使用发货计划数据预测
    - order: ARIMA模型的阶数，形如 (p, d, q)。
    - forecast_steps: 预测的时间步数量。

    返回：
    - 预测结果，维度为（forecast_steps，n_samples）。
    """

    # 存储预测结果
    predictions = []

    # 对每个时间序列进行训练和预测
    for i in range(data.shape[1]):
        # 训练ARIMA模型
        if use_plan == 'False':
            model = ARIMA(data[:, i], order=order)
            model_fit = model.fit()
            forecast = model_fit.get_forecast(steps=forecast_steps)
            forecast_mean = forecast.predicted_mean
            predictions.append(forecast_mean)

        else:
            model = ARIMA(data[:, i], exog=plan[:, i], order=order)
            model_fit = model.fit()
            forecast = model_fit.get_forecast(steps=forecast_steps, exog=f_plan[:, i])
            forecast_mean = forecast.predicted_mean
            predictions.append(forecast_mean)

        # 进行未来预测  # forecast = model_fit.get_forecast(steps=forecast_steps)

        # 获取预测结果  # forecast_mean = forecast.predicted_mean  # predictions.append(forecast_mean)

    # 将预测结果转换为数组，维度为（forecast_steps，n_samples）
    predictions = np.array(predictions).T
    # predictions = pd.DataFrame(predictions)
    predictions[predictions < 0] = 0
    return predictions


def forecast_with_autoarima(data, plan, f_plan, use_plan, order, s, d_range, D_range, p_range, q_range, P_range,
                            Q_range, forecast_steps):
    """
    使用AutoARIMA模型进行预测

    参数：
    - df: 输入的DataFrame，包含时间序列数据，形状为(N, F)，N代表时间步，F代表不同的时间序列
    - plan: 输入的发货计划序列数据，维度为（n_timesteps, n_samples），要求与data维度一样。
    - use_plan: 是否使用发货计划数据预测
    - s: 季节性周期长度
    - d_range: 差分阶数的范围，例如(0, 1)表示可选0或1阶差分
    - D_range: 季节性差分阶数的范围，例如(0, 1)表示可选0或1阶季节性差分
    - p_range: 自回归阶数的范围，例如(0, 2, 4)表示可选0、1或2阶自回归
    - q_range: 移动平均阶数的范围，例如(0, 2, 4)表示可选0、1或2阶移动平均
    - P_range: 季节性自回归阶数的范围，例如range(2)表示可选0或1阶季节性自回归
    - Q_range: 季节性移动平均阶数的范围，例如range(2)表示可选0或1阶季节性移动平均
    - forecast_steps: 预测的时间步数

    返回值：
    - fc: 预测结果，形状为(F, forecast_steps)，每列包含对应时间序列的未来预测值
    """
    from cuml.tsa.arima import ARIMA
    from cuml.tsa.auto_arima import AutoARIMA

    if use_plan == 'False':
        model = AutoARIMA(data)
        # 搜索最佳模型参数
        model.search(s=s, d=d_range, D=D_range, p=p_range, q=q_range, P=P_range, Q=Q_range, method="css", truncate=20)
        model.fit(method="css-ml")
        fc = model.forecast(forecast_steps)
    else:  # 如果要用发货计划数据, 则不能使用AutoArima
        # model = AutoARIMA()
        # model.fit(data, exogenous=plan)
        # forecast = model.predict(n_periods=forecast_steps, exogenous= f_plan)
        # model = ARIMA(data, exog=)
        noise_level = 1e-6
        data += np.random.uniform(0, noise_level, size=(data.shape[0], data.shape[1]))
        model = ARIMA(data, exog=plan, order=(1, 0, 0))  # exog=plan
        model.fit()
        fc = model.forecast(forecast_steps, exog=f_plan)  # exog=f_plan
    # 拟合模型
    # model.fit(method="css-ml")
    # 进行未来预测
    # fc = model.forecast(forecast_steps)
    return fc


def forecast(data, plan, f_plan, forecast_steps=3, use_plan=False, use_gpu=False):
    """
    对多个时间序列进行预测，可选择使用ARIMA模型或AutoARIMA模型。

    参数：
    - data: 输入的时间序列数据，维度为（n_timesteps, n_samples）。
    - plan: 输入的发货计划序列数据，维度为（n_timesteps, n_samples），要求与data维度一样。
    - order: ARIMA模型的阶数，形如 (p, d, q)。
    - forecast_steps: 预测的时间步数量。
    - use_plan: 是否使用发货计划数据预测
    - use_gpu: 是否使用AutoARIMA模型。

    返回：
    - 预测结果，维度为（forecast_steps，n_samples）。
    """
    if use_gpu == 'True':
        order = (1, 0, 0)
        forecast = forecast_with_autoarima(data, plan, f_plan, use_plan, order, s=2, d_range=(0, 5), D_range=(0, 5),
                                           p_range=(0, 5), q_range=(0, 5), P_range=(0, 5), Q_range=(0, 5),
                                           forecast_steps=forecast_steps)
        print(forecast)  # 输出格式为 (forecast_steps, num_variables)  forecast_steps为未来预测的步数，num_variables为时间序列的个数
        return forecast
    else:
        order = (1, 0, 0)  # 无GPU情况下可以调整的参数,备选（1,0,0）（0,1,1）
        return arima_forecast(data, plan, f_plan=f_plan, use_plan=use_plan, order=order, forecast_steps=forecast_steps)
